from flask import Flask, request, jsonify
from flask_sqlalchemy import SQLAlchemy
from flask_socketio import SocketIO, join_room, leave_room, emit
from datetime import datetime
import hashlib
import os
from flask_cors import CORS

# 初始化Flask应用
app = Flask(__name__)
app.config['SECRET_KEY'] = os.urandom(24)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///chat.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

# 初始化扩展
db = SQLAlchemy(app)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode="threading")
CORS(app)

# 数据库模型
class User(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(50), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

    def set_password(self, password):
        self.password_hash = hashlib.sha256(password.encode()).hexdigest()

    def check_password(self, password):
        return self.password_hash == hashlib.sha256(password.encode()).hexdigest()

class Message(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    sender_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    receiver_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=True)  # None表示公共消息
    content = db.Column(db.Text, nullable=False)
    timestamp = db.Column(db.DateTime, default=datetime.utcnow)
    
    sender = db.relationship('User', foreign_keys=[sender_id], backref='sent_messages')
    receiver = db.relationship('User', foreign_keys=[receiver_id], backref='received_messages')

# 在线用户管理 - 修改为存储房间信息
online_users = {}  # {user_id: (username, sid, room)}
user_rooms = {}    # {user_id: room_name} - 用于跟踪用户所在房间

# API路由
@app.route('/api/register', methods=['POST'])
def register():
    data = request.json
    username = data.get('username')
    password = data.get('password')
    
    if not username or not password:
        return jsonify({'success': False, 'message': '用户名和密码不能为空'}), 400
    
    if User.query.filter_by(username=username).first():
        return jsonify({'success': False, 'message': '用户名已存在'}), 400
    
    user = User(username=username)
    user.set_password(password)
    
    try:
        db.session.add(user)
        db.session.commit()
        return jsonify({'success': True, 'message': '注册成功'})
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': f'注册失败: {str(e)}'}), 500

@app.route('/api/login', methods=['POST'])
def login():
    data = request.json
    username = data.get('username')
    password = data.get('password')
    
    user = User.query.filter_by(username=username).first()
    
    if not user or not user.check_password(password):
        return jsonify({'success': False, 'message': '用户名或密码错误'}), 401
    
    return jsonify({
        'success': True, 
        'message': '登录成功',
        'user_id': user.id,
        'username': user.username
    })

# WebSocket事件处理
@socketio.on('connect')
def handle_connect():
    print('Client connected:', request.sid)

@socketio.on('disconnect')
def handle_disconnect():
    print('Client disconnected:', request.sid)
    # 从在线用户中移除
    user_to_remove = None
    for user_id, (username, sid, room) in online_users.items():
        if sid == request.sid:
            user_to_remove = user_id
            break
    
    if user_to_remove:
        username = online_users[user_to_remove][0]
        room = online_users[user_to_remove][2]
        del online_users[user_to_remove]
        if user_to_remove in user_rooms:
            del user_rooms[user_to_remove]
        # 广播用户离线
        emit('user_offline', {'user_id': user_to_remove, 'username': username}, room=room)

@socketio.on('login')
def handle_login(data):
    user_id = data.get('user_id')
    username = data.get('username')
    
    if user_id and username:
        # 为每个用户创建一个私人房间
        room = f'user_{user_id}'
        join_room(room)
        
        # 更新在线用户信息
        online_users[user_id] = (username, request.sid, room)
        user_rooms[user_id] = room
        
        # 加入公共房间
        public_room = 'public_chat'
        join_room(public_room)
        
        # 广播新用户上线到公共房间
        emit('user_online', {'user_id': user_id, 'username': username}, room=public_room)
        
        # 发送当前在线用户列表给新登录用户
        current_online = [{'user_id': uid, 'username': uname} for uid, (uname, _, _) in online_users.items()]
        emit('online_users', current_online)

@socketio.on('public_message')
def handle_public_message(data):
    sender_id = data.get('sender_id')
    content = data.get('content')
    
    if sender_id not in online_users or not content:
        return
    
    sender_username = online_users[sender_id][0]
    public_room = 'public_chat'
    
    # 保存消息到数据库
    message = Message(
        sender_id=sender_id,
        receiver_id=None,  # 公共消息
        content=content
    )
    
    try:
        db.session.add(message)
        db.session.commit()
        
        # 广播公共消息到公共房间
        emit('public_message', {
            'message_id': message.id,
            'sender_id': sender_id,
            'sender_username': sender_username,
            'content': content,
            'timestamp': message.timestamp.isoformat()
        }, room=public_room)
    except Exception as e:
        db.session.rollback()
        print(f"保存消息失败: {str(e)}")

@socketio.on('private_message')
def handle_private_message(data):
    sender_id = data.get('sender_id')
    receiver_id = data.get('receiver_id')
    content = data.get('content')
    
    if sender_id not in online_users or receiver_id not in online_users or not content:
        return
    
    sender_username = online_users[sender_id][0]
    receiver_room = online_users[receiver_id][2]  # 获取接收者的私人房间
    
    # 保存消息到数据库
    message = Message(
        sender_id=sender_id,
        receiver_id=receiver_id,
        content=content
    )
    
    try:
        db.session.add(message)
        db.session.commit()
        
        # 只发送给接收者的私人房间
        emit('private_message', {
            'message_id': message.id,
            'sender_id': sender_id,
            'sender_username': sender_username,
            'content': content,
            'timestamp': message.timestamp.isoformat()
        }, room=receiver_room)
        
        # 发送给自己（确认消息已发送）
        emit('private_message_sent', {
            'message_id': message.id,
            'receiver_id': receiver_id,
            'receiver_username': online_users[receiver_id][0],
            'content': content,
            'timestamp': message.timestamp.isoformat()
        })
    except Exception as e:
        db.session.rollback()
        print(f"保存私信失败: {str(e)}")

# 应用入口
if __name__ == '__main__':
    # 创建数据库表
    with app.app_context():
        db.create_all()
    # 启动服务
    socketio.run(app, debug=True, host='0.0.0.0', port=5000)